# 流式输出
# https://help.aliyun.com/zh/model-studio/stream


from openai import OpenAI, APIError

# 1. 准备客户端
try:
    client = OpenAI(
        api_key = "sk-1341cd7d8e38416ebc0a0e698a622896", #os.getenv("DASHCOPE_API_KEY")
        base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
except KeyError:
    raise ValueError("正确设置api-key")

def stream_task():
    try:
        # 发起流式请求
        completion = client.chat.completions.create(
            model = "qwen-plus",
            messages = [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "请介绍一下自己"}
            ],
            stream= True, # 流式
            # 目的：在最后一个chunk中获取本次请求的Token用量。
            stream_options = {"include_usage": True}
        )

        # 处理流式响应
        content_parts = []
        print("AI:", end="", flush=True)
        for chunk in completion:
            # 最后一个chunk不包含choices，但包含usage
            if chunk.choices:
                # 关键：delta.content可能为None，使用`or ""`避免拼接错误
                content = chunk.choices[0].delta.content or ""
                print(content, end ="", flush=True)
                content_parts.append(content)
            elif chunk.usage:
                # 请求结束
                print("\n--- 请求用量 ---")
                print(f"输入Tokens：{chunk.usage.prompt_tokens}")
                print(f"输出Tokens：{chunk.usage.completion_tokens}")
                print(f"总计Tokens：{chunk.usage.total_tokens}")
        full_response = "".join(content_parts)
    except APIError as e:
        print(f"API请求失败：{e}")
    except Exception as e:
        print(f"错误信息：{e}")
        print("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code")

def main():
    stream_task()


if __name__ == "__main__":
    main()
